#include <iostream>
#include <fstream>
#include <vector>
#include <map>
#include <string>
#include <array>
#include <pcl/io/pcd_io.h>
#include <pcl/point_types.h>
#include <pcl/visualization/pcl_visualizer.h>
#include <onnxruntime_cxx_api.h>
#include <Eigen/Core>

// ========== 体素参数 ==========
constexpr float voxel_size_x = 0.16f, voxel_size_y = 0.16f, voxel_size_z = 4.0f;
constexpr float x_min = 0.0f, x_max = 69.12f;
constexpr float y_min = -39.68f, y_max = 39.68f;
constexpr float z_min = -3.0f, z_max = 1.0f;
constexpr int max_points_per_voxel = 32;
constexpr int max_voxels = 16000;
constexpr int input_dim = 4; // [x, y, z, intensity]

// ========== 1. 读取KITTI bin ==========
void readKittiBin(const std::string& bin_path, pcl::PointCloud<pcl::PointXYZI>::Ptr cloud) {
    std::ifstream input(bin_path, std::ios::binary);
    if (!input) { std::cerr << "open file fail!\n";  return; }
    float data[4];
    while (input.read(reinterpret_cast<char*>(data), sizeof(data))) {
        pcl::PointXYZI pt; pt.x = data[0]; pt.y = data[1]; pt.z = data[2]; pt.intensity = data[3];
        cloud->push_back(pt);
    }
    input.close();
}

// ========== 2. 体素化 ==========
void voxelization(const pcl::PointCloud<pcl::PointXYZI>::Ptr& cloud,
                  std::vector<float> &voxels, std::vector<int> &coors,
                  std::vector<int> &num_points, int& num_voxels)
{
    struct Key {
        int z, y, x;
        bool operator<(const Key&o) const {
            if (z != o.z) return z < o.z;
            if (y != o.y) return y < o.y;
            return x < o.x;
        }
    };

    std::map<Key, std::vector<std::array<float, input_dim>>> voxel_map;

    for (const auto& pt : cloud->points) {
        if (!(pt.x >= x_min && pt.x < x_max &&
              pt.y >= y_min && pt.y < y_max &&
              pt.z >= z_min && pt.z < z_max))
            continue;
        int x_idx = int((pt.x - x_min) / voxel_size_x);
        int y_idx = int((pt.y - y_min) / voxel_size_y);
        int z_idx = int((pt.z - z_min) / voxel_size_z);
        Key key{z_idx, y_idx, x_idx};
        auto& voxel = voxel_map[key];
        if ((int)voxel.size() < max_points_per_voxel)
            voxel.push_back({pt.x, pt.y, pt.z, pt.intensity});
        if ((int)voxel_map.size() >= max_voxels) break;
    }
    num_voxels = std::min((int)voxel_map.size(), max_voxels);

    // 输出buffer+padding
    voxels.assign(max_voxels * max_points_per_voxel * input_dim, 0.f);
    coors.assign(max_voxels * 4, 0); // [batch_idx, z, y, x]
    num_points.assign(max_voxels, 0);

    int idx = 0;
    for (auto it = voxel_map.begin(); idx < num_voxels && it != voxel_map.end(); ++it, ++idx) {
        const auto& key = it->first;
        const auto& points = it->second;
        for (size_t p = 0; p < points.size(); ++p) {
            for (int c = 0; c < input_dim; ++c)
                voxels[idx * max_points_per_voxel * input_dim + p * input_dim + c] = points[p][c];
        }
        num_points[idx] = points.size();
        coors[idx * 4 + 0] = 0;  // batch idx
        coors[idx * 4 + 1] = key.z;
        coors[idx * 4 + 2] = key.y;
        coors[idx * 4 + 3] = key.x;
    }
}

// ========== 3. 可视化包围框 ==========
void draw_boxes(pcl::visualization::PCLVisualizer::Ptr vis, const std::vector<std::vector<float>>& boxes,
                const std::vector<float>& scores, const std::vector<int>& labels,
                const std::vector<std::string>& cls_names, float score_thr=0.3)
{
    for(size_t i=0; i<boxes.size(); ++i) {
        if(scores[i]<score_thr) continue;
        float x = boxes[i][0], y = boxes[i][1], z = boxes[i][2];
        float dx = boxes[i][3], dy = boxes[i][4], dz = boxes[i][5], yaw = boxes[i][6];
        Eigen::Vector3f translation(x, y, z);
        Eigen::Quaternionf quat(Eigen::AngleAxisf(yaw, Eigen::Vector3f::UnitZ()));
        std::string id = "box_"+std::to_string(i);
        double r=0,g=0,b=1;
        if(!cls_names.empty()) {
            if(labels[i]==0){r=1;g=0;b=0;} // Pedestrian
            if(labels[i]==1){r=0;g=1;b=0;} // Cyclist
            if(labels[i]==2){r=1;g=1;b=0;} // Car
        }
        vis->addCube(translation, quat, dx, dy, dz, id);
        vis->setShapeRenderingProperties(pcl::visualization::PCL_VISUALIZER_COLOR, r,g,b, id);
        vis->setShapeRenderingProperties(pcl::visualization::PCL_VISUALIZER_LINE_WIDTH, 2.0, id);
        vis->addText3D(cls_names[labels[i]]+" "+std::to_string(scores[i]).substr(0,4), pcl::PointXYZ(x,y,z+dz/2+1),1, r,g,b, id+"_label");
    }
}

// ========== 4. 主流程 ==========
int main() {
    std::string kitti_bin_file = "/home/chenhao/kitti3d_segmentation_benchmark/data/bin/000000.bin";
    std::string onnx_path = "/home/chenhao/kitti3d_segmentation_benchmark/onnx/end2end.onnx";
    std::vector<std::string> cls_names = { "Pedestrian","Cyclist","Car"};

    // 1. 点云读取
    pcl::PointCloud<pcl::PointXYZI>::Ptr cloud(new pcl::PointCloud<pcl::PointXYZI>());
    readKittiBin(kitti_bin_file, cloud);
    std::cout << "Load points: " << cloud->size() << std::endl;
    if (cloud->empty()) return -1;

    // 2. 体素化
    std::vector<float> voxels;
    std::vector<int> coors;
    std::vector<int> num_points;
    int num_voxels = 0;
    voxelization(cloud, voxels, coors, num_points, num_voxels);
    std::cout << "[Voxelization] num_voxels = " << num_voxels << std::endl;

    // 3. ONNX推理
    Ort::Env env(ORT_LOGGING_LEVEL_WARNING, "det3d");
    Ort::SessionOptions session_options;
    Ort::Session session(env, onnx_path.c_str(), session_options);

    // 获取输入输出名
    std::vector<std::string> input_names_raw = session.GetInputNames();
    std::vector<const char*> input_names;
    for (const auto& s : input_names_raw) input_names.push_back(s.c_str());

    std::vector<std::string> output_names_raw = session.GetOutputNames();
    std::vector<const char*> output_names;
    for (const auto& s : output_names_raw) output_names.push_back(s.c_str());

    std::vector<int64_t> voxels_shape{num_voxels, max_points_per_voxel, input_dim};
    std::vector<int64_t> num_points_shape{num_voxels};
    std::vector<int64_t> coors_shape{num_voxels, 4};

    Ort::MemoryInfo memInfo = Ort::MemoryInfo::CreateCpu(OrtArenaAllocator, OrtMemTypeDefault);
    Ort::Value voxels_tensor = Ort::Value::CreateTensor<float>(
        memInfo, voxels.data(), voxels.size(), voxels_shape.data(), voxels_shape.size());
    Ort::Value num_points_tensor = Ort::Value::CreateTensor<int32_t>(
        memInfo, (int32_t*)num_points.data(), num_points.size(), num_points_shape.data(), num_points_shape.size());
    Ort::Value coors_tensor = Ort::Value::CreateTensor<int32_t>(
        memInfo, (int32_t*)coors.data(), coors.size(), coors_shape.data(), coors_shape.size());
    std::vector<Ort::Value> input_tensors;
    input_tensors.push_back(std::move(voxels_tensor));
    input_tensors.push_back(std::move(num_points_tensor));
    input_tensors.push_back(std::move(coors_tensor));

    // 推理
    auto output_tensors = session.Run(Ort::RunOptions{nullptr},
        input_names.data(), input_tensors.data(), input_tensors.size(),
        output_names.data(), output_names.size());

    // 4. 后处理
    auto& out_boxes = output_tensors[0];
    auto& out_scores = output_tensors[1];
    auto& out_labels = output_tensors[2];
    size_t num_det = out_scores.GetTensorTypeAndShapeInfo().GetShape()[0];
    const float* boxes_ptr = out_boxes.GetTensorData<float>();
    const float* scores_ptr = out_scores.GetTensorData<float>();
    const int64_t* labels_ptr = out_labels.GetTensorData<int64_t>();

    std::vector<std::vector<float>> det_boxes;
    std::vector<float> det_scores;
    std::vector<int> det_labels;
    for(size_t i=0; i<num_det; ++i) {
        std::vector<float> box(boxes_ptr+i*7, boxes_ptr+(i+1)*7);
        det_boxes.push_back(box);
        det_scores.push_back(scores_ptr[i]);
        det_labels.push_back(labels_ptr[i]);

        // 详细打印每个检测框
        std::cout << "box[" << i << "]: [";
        for(int j=0; j<7; ++j)
            std::cout << box[j] << (j==6 ? "" : ", ");
        std::cout << "]"
                << ", label: " << det_labels.back()
                << ", score: " << det_scores.back()
                << std::endl;
    }
    std::cout << "[Detect boxes]: " << det_boxes.size() << std::endl;

    // 5. 可视化
    pcl::visualization::PCLVisualizer::Ptr viewer(new pcl::visualization::PCLVisualizer("Det3D Viewer"));
    viewer->addPointCloud<pcl::PointXYZI>(cloud, "cloud0");
    draw_boxes(viewer, det_boxes, det_scores, det_labels, cls_names, 0.3f);

    while (!viewer->wasStopped()) viewer->spinOnce();

    return 0;
}